#include "wake_word_detect.h"
#include "application.h"

#include <esp_log.h>
#include <model_path.h>
#include <arpa/inet.h>
#include <sstream>

#include "freertos/FreeRTOS.h"
#include "freertos/queue.h"
#include "freertos/semphr.h"

#include "esp_system.h"
//#include "esp_wifi.h"
//#include "esp_event.h"
//#include "esp_log.h"
//#include "esp_mac.h"
//#include "nvs_flash.h"
//#include "sys/socket.h"
//#include "lwip/err.h"
//#include "lwip/sys.h"
//#include "netdb.h"
//#include "I2Sto.h"
#include <string.h>
//#include "wifito.h"
// #include "wifito.h"

#include "esp_mn_models.h"
#include "esp_mn_iface.h"
#include "esp_wn_iface.h"
#include "esp_wn_models.h"
#include "esp_afe_sr_models.h"
#include "esp_afe_sr_iface.h"

//#include "driver/i2s.h"
#include "esp_mn_speech_commands.h"

#define DETECTION_RUNNING_EVENT 1

static const char* TAG = "WakeWordDetect";

WakeWordDetect::WakeWordDetect()
    : afe_data_(nullptr),
      wake_word_pcm_(),
      wake_word_opus_() {

    event_group_ = xEventGroupCreate();
}

WakeWordDetect::~WakeWordDetect() {
    if (afe_data_ != nullptr) {
        afe_iface_->destroy(afe_data_);
    }

    if (wake_word_encode_task_stack_ != nullptr) {
        heap_caps_free(wake_word_encode_task_stack_);
    }

    vEventGroupDelete(event_group_);
}

typedef struct {
    wakenet_state_t     wakenet_mode;
    esp_mn_state_t      state;
    int                 command_id;
} sr_result_t;

static QueueHandle_t            g_result_que    = NULL;
int detect_flag = 0;
static esp_afe_sr_iface_t *afe_handle = NULL;
static volatile int task_flag = 0;
static model_iface_data_t       *model_data     = NULL;
static const esp_mn_iface_t     *multinet       = NULL;

const char *cmd_phoneme[3] = {
    "ni hao mi tang",
    "da kai kong qi jing hua qi",
    "guan bi kong qi jing hua qi"  //修改
    // "da kai tai deng",
    // "guan bi tai deng",
    // "tai deng tiao liang",
    // "tai deng tiao an",
    // "da kai deng dai",
    // "guan bi deng dai",
    // "bo fang yin yue",
    // "ting zhi bo fang",
    // "da kai shi jian",
    //"da kai ri li"
};

void print_model_info(srmodel_list_t *models) {
    if (models == NULL) {
        printf("Model list is NULL!\n");
        return;
    }

    printf("Model list info:\n");
    printf("  Model count: %d\n", models->num);
    
    // if (models->num > 0 && models->model_list != NULL) {
    //     printf("  Model details:\n");
    //     for (int i = 0; i < models->model_num; i++) {
    //         srmodel_t *model = models->model_list[i];
    //         if (model != NULL) {
    //             printf("    Model %d:\n", i);
    //             printf("      Name: %s\n", model->model_name ? model->model_name : "NULL");
    //             printf("      Type: %d\n", model->model_type);
    //             printf("      Language: %d\n", model->model_language);
    //             printf("      Size: %d bytes\n", model->model_size);
    //             // 可根据需要添加更多字段
    //         } else {
    //             printf("    Model %d: NULL\n", i);
    //         }
    //     }
    // } else {
    //     printf("  No models found in list.\n");
    // }
}

int WakeWordDetect::Initialize(int channels, bool reference) {
    channels_ = channels;
    reference_ = reference;
    int ref_num = reference_ ? 1 : 0;

    srmodel_list_t *models = esp_srmodel_init("model");

    print_model_info(models);

    std::string input_format;
    for (int i = 0; i < channels_ - ref_num; i++) {
        input_format.push_back('M');
    }
    for (int i = 0; i < ref_num; i++) {
        input_format.push_back('R');
    }
    afe_config_t* afe_config = afe_config_init("M", models, AFE_TYPE_SR, AFE_MODE_HIGH_PERF);
    afe_config->aec_init = reference_;
    afe_config->aec_mode = AEC_MODE_SR_HIGH_PERF;
    afe_config->afe_perferred_core = 1;
    afe_config->afe_perferred_priority = 1;
    afe_config->memory_alloc_mode = AFE_MEMORY_ALLOC_MORE_PSRAM;
    
    afe_iface_ = esp_afe_handle_from_config(afe_config);
    afe_data_ = afe_iface_->create_from_config(afe_config);

    char *mn_name = esp_srmodel_filter(models, ESP_MN_CHINESE, NULL);
    if (NULL == mn_name) {
        printf("No multinet model found");
        return ESP_FAIL;
    }
    multinet = esp_mn_handle_from_name(mn_name);
    model_data = multinet->create(mn_name, 5760);//设置唤醒超时时间
    //printf( "load multinet:%s", mn_name);
    esp_mn_commands_clear();//清除唤醒指令
    for (int i = 0; i < sizeof(cmd_phoneme) / sizeof(cmd_phoneme[0]); i++) {
        esp_mn_commands_add(i, (char *)cmd_phoneme[i]);//逐个将唤醒指令放入
    }
    esp_mn_commands_update();//更新命令词列表
    esp_mn_commands_print();
    multinet->print_active_speech_commands(model_data);//输出目前激活的命令词

    for (int i = 0; i < models->num; i++) {
        ESP_LOGI(TAG, "Model %d: %s", i, models->model_name[i]);
        if (strstr(models->model_name[i], ESP_WN_PREFIX) != NULL) {
            wakenet_model_ = models->model_name[i];
            auto words = esp_srmodel_get_wake_words(models, wakenet_model_);
            // split by ";" to get all wake words
            std::stringstream ss(words);
            std::string word;
            while (std::getline(ss, word, ';')) {
                wake_words_.push_back(word);
            }
        }
    }

    detect_flag=1;

    xTaskCreate([](void* arg) {
        auto this_ = (WakeWordDetect*)arg;
        this_->AudioDetectionTask();
        vTaskDelete(NULL);
    }, "audio_detection", 4096, this, 3, nullptr);
    return 0;
}

void WakeWordDetect::OnWakeWordDetected(std::function<void(const std::string& wake_word)> callback) {
    wake_word_detected_callback_ = callback;
}

void WakeWordDetect::StartDetection() {
    xEventGroupSetBits(event_group_, DETECTION_RUNNING_EVENT);
}

void WakeWordDetect::StopDetection() {
    xEventGroupClearBits(event_group_, DETECTION_RUNNING_EVENT);
    afe_iface_->reset_buffer(afe_data_);
}

bool WakeWordDetect::IsDetectionRunning() {
    return xEventGroupGetBits(event_group_) & DETECTION_RUNNING_EVENT;
}

void WakeWordDetect::Feed(const std::vector<int16_t>& data) {
    input_buffer_.insert(input_buffer_.end(), data.begin(), data.end());

    auto feed_size = afe_iface_->get_feed_chunksize(afe_data_) * channels_;
    while (input_buffer_.size() >= feed_size) {
        afe_iface_->feed(afe_data_, input_buffer_.data());
        input_buffer_.erase(input_buffer_.begin(), input_buffer_.begin() + feed_size);
    }
}

void WakeWordDetect::AudioDetectionTask() {
    auto fetch_size = afe_iface_->get_fetch_chunksize(afe_data_);
    auto feed_size = afe_iface_->get_feed_chunksize(afe_data_);
    ESP_LOGI(TAG, "Audio detection task started, feed size: %d fetch size: %d",
        feed_size, fetch_size);

    while (true) {
        xEventGroupWaitBits(event_group_, DETECTION_RUNNING_EVENT, pdFALSE, pdTRUE, portMAX_DELAY);

        auto res = afe_iface_->fetch_with_delay(afe_data_, portMAX_DELAY);
        if (res == nullptr || res->ret_value == ESP_FAIL) {
            continue;;
        }
        // if(model_data==nullptr)
        //     continue;
        esp_mn_state_t mn_state = ESP_MN_STATE_DETECTING;

            mn_state = multinet->detect(model_data, res->data);

            // if (ESP_MN_STATE_DETECTING == mn_state) {
            //     continue;
            // }

            // if (ESP_MN_STATE_TIMEOUT == mn_state) {//超时
            //     ESP_LOGW(TAG, "Time out");
            //     sr_result_t result = {
            //         .wakenet_mode = WAKENET_NO_DETECT,
            //         .state = mn_state,
            //         .command_id = 0,
            //     };
            //     xQueueSend(g_result_que, &result, 10);
            //     afe_handle->enable_wakenet(afe_data);
            //     detect_flag = false;
            //     continue;
            // }

            if (ESP_MN_STATE_DETECTED == mn_state) {//得到指令
                StopDetection();
            //last_detected_wake_word_ = wake_words_[res->wake_word_index - 1];

            if (wake_word_detected_callback_) {
                wake_word_detected_callback_(last_detected_wake_word_);
            }

                // esp_mn_results_t *mn_result = multinet->get_results(model_data);
                // for (int i = 0; i < mn_result->num; i++) {
                //     ESP_LOGI(TAG, "TOP %d, command_id: %d, phrase_id: %d, prob: %f",
                //             i + 1, mn_result->command_id[i], mn_result->phrase_id[i], mn_result->prob[i]);
                // }

                // int sr_command_id = mn_result->command_id[0];
                // ESP_LOGI(TAG, "Deteted command : %d", sr_command_id);
                // sr_result_t result = {
                //     .wakenet_mode = WAKENET_NO_DETECT,
                //     .state = mn_state,
                //     .command_id = sr_command_id,
                // };
                // xQueueSend(g_result_que, &result, 10);
            }

        // Store the wake word data for voice recognition, like who is speaking
        StoreWakeWordData((uint16_t*)res->data, res->data_size / sizeof(uint16_t));

        if (res->wakeup_state == WAKENET_DETECTED) {
            StopDetection();
            last_detected_wake_word_ = wake_words_[res->wake_word_index - 1];

            if (wake_word_detected_callback_) {
                wake_word_detected_callback_(last_detected_wake_word_);
            }
        }
    }
}

void WakeWordDetect::StoreWakeWordData(uint16_t* data, size_t samples) {
    // store audio data to wake_word_pcm_
    wake_word_pcm_.emplace_back(std::vector<int16_t>(data, data + samples));
    // keep about 2 seconds of data, detect duration is 32ms (sample_rate == 16000, chunksize == 512)
    while (wake_word_pcm_.size() > 2000 / 32) {
        wake_word_pcm_.pop_front();
    }
}

void WakeWordDetect::EncodeWakeWordData() {
    wake_word_opus_.clear();
    if (wake_word_encode_task_stack_ == nullptr) {
        wake_word_encode_task_stack_ = (StackType_t*)heap_caps_malloc(4096 * 8, MALLOC_CAP_SPIRAM);
    }
    wake_word_encode_task_ = xTaskCreateStatic([](void* arg) {
        auto this_ = (WakeWordDetect*)arg;
        {
            auto start_time = esp_timer_get_time();
            auto encoder = std::make_unique<OpusEncoderWrapper>(16000, 1, OPUS_FRAME_DURATION_MS);
            encoder->SetComplexity(0); // 0 is the fastest

            for (auto& pcm: this_->wake_word_pcm_) {
                encoder->Encode(std::move(pcm), [this_](std::vector<uint8_t>&& opus) {
                    std::lock_guard<std::mutex> lock(this_->wake_word_mutex_);
                    this_->wake_word_opus_.emplace_back(std::move(opus));
                    this_->wake_word_cv_.notify_all();
                });
            }
            this_->wake_word_pcm_.clear();

            auto end_time = esp_timer_get_time();
            ESP_LOGI(TAG, "Encode wake word opus %zu packets in %lld ms",
                this_->wake_word_opus_.size(), (end_time - start_time) / 1000);

            std::lock_guard<std::mutex> lock(this_->wake_word_mutex_);
            this_->wake_word_opus_.push_back(std::vector<uint8_t>());
            this_->wake_word_cv_.notify_all();
        }
        vTaskDelete(NULL);
    }, "encode_detect_packets", 4096 * 8, this, 2, wake_word_encode_task_stack_, &wake_word_encode_task_buffer_);
}

bool WakeWordDetect::GetWakeWordOpus(std::vector<uint8_t>& opus) {
    std::unique_lock<std::mutex> lock(wake_word_mutex_);
    wake_word_cv_.wait(lock, [this]() {
        return !wake_word_opus_.empty();
    });
    opus.swap(wake_word_opus_.front());
    wake_word_opus_.pop_front();
    return !opus.empty();
}
